# 负责与配置多类大模型
# 本文使用的大模型均用ollama进行本质部署


class LLMModelConfig:
    def __init__(self, api_key, base_url):
        self.api_key = api_key
        self.base_url = base_url

    def __repr__(self):
        # 隐藏完整的 API 密钥，只显示前 6 个字符
        masked_api_key = self.api_key[:6] + "*" * (len(self.api_key) - 6)
        return f"LLMModelConfig(api_key='{masked_api_key}', base_url='{self.base_url}')"

class LLMAPIConfig:
    # LLM 配置
    MODELS = {
        "deepseek-r1:1.5b": LLMModelConfig(
            api_key="ollama",
            base_url="http://127.0.0.1:11434/v1"
        ),
        "qwen2.5-coder": LLMModelConfig(
            api_key="ollama",
            base_url="http://127.0.0.1:11434/v1"
        ),
        "llama3.1:8b": LLMModelConfig(
            api_key="ollama",
            base_url="http://127.0.0.1:11434/v1"
        ),
    }

    # 具体使用的模型设置
    TASK_MODELS = {
        "planner": "llama3.1:8b",
        "coder": "qwen2.5-coder",
        "summarizer": "llama3.1:8b"
    }
    @classmethod
    def get_model_config(cls, model_name):
        return cls.MODELS.get(model_name)

    @classmethod
    def get_task_model(cls, task):
        model_name = cls.TASK_MODELS.get(task)
        return cls.get_model_config(model_name)

if __name__ == '__main__':
    config = LLMAPIConfig()
    print("Model config for qwen2.5-coder:")
    print(config.get_model_config("qwen2.5-coder"))
    print("\nModel config for planner task:")
    print(config.get_task_model("planner"))